# -*- coding: utf-8 -*-
"""
航班查询

可以查询 国内、国际 机票

@Author: rslai
@Email: netlrs@163.com
@Date: 2025-09-17
"""
from base.tool_registry import mcp_tool
import requests
from tools.flight_fliggy_data_parse import FlightFliggyDataParse
from utils.city_utils import CityUtils
from utils.file_utils import write_to_file
from utils.ai_utils import AIUtils
from tools.flight_conf import FlightConf
from utils.uuid_utils import generate_timestamp_id
import time
import random
import pandas as pd
import os

flight_optimizer_prompt = """
## 角色
作为一名专业的机票推荐专员，请根据以下航班数据表格，严格按照以下筛选条件选出3个最优航班：

## 筛选条件（按优先级排序）：
1. 价格优先 - 选择价格最低的航班
2. 飞行时长 - 在价格相同的情况下，选择飞行时长更短的航班

## 输出要求：
1. 仅输出3个最优航班的ID
2. ID之间用英文逗号分隔
3. 不要包含任何其他文字或解释
4. 不要显示思考过程或分析步骤

## 数据说明：
1. 表格列包括：ID、出发机场、到达机场、出发时间、到达时间、飞行时长、航空公司、航班号、机型、中转、航段数量、中转城市、中转时间、价格


## 航班数据
=== 
{data}
===
"""



@mcp_tool()
def search_flights_between_countries(
    departure_country: str="中国", 
    arrival_country: str="美国", 
    departure_date: str="2026/02/10",
    max_duration: int=0,
    max_total_price: int=0,
    sort_by: list[str]=[],
    limit: int=0
):
    """
    可以查询两个国家之间的航班信息。
    
    可以查询指定日期的从"出发国家"到"到达国家"的机票信息。返回机票的出发到达城市、价格、飞行时长等信息的结构化数据。
    
    Args:
        departure_country (str): 出发国家
        arrival_country (str): 到达国家
        departure_date (str): 出发日期，格式为 YYYY-MM-DD
        max_duration (int): 最大飞行时长，单位：小时，默认值为 24小时
        max_total_price (int): 机票最大总价，单位：元，默认值为 5000元
        sort_by (list[str]): 排序规则字段，默认值为 ['total_price', 'duration'] 默认先按总价格排序，再按持续时间排序
        limit (int): 返回最优的记录条数，默认为20条

    Returns:
        str: 机票查询结果
    """
    conf = FlightConf() # 配置信息
    conf.work_dir = f"{os.getcwd()}/temp/{generate_timestamp_id()}" # 设置工作目录

    if departure_country.strip(): 
        conf.departure_country = departure_country
    else:
        return False, "出发国家 必填!"

    if arrival_country.strip(): 
        conf.arrival_country = arrival_country
    else:
        return False, "到达国家 必填!"

    if departure_date.strip(): 
        conf.departure_date = departure_date
    else:
        return False, "出发日期 必填!"

    if max_duration > 0: conf.max_duration = max_duration
    if max_total_price > 0: conf.max_total_price = max_total_price
    if len(sort_by) > 0: conf.sort_by = sort_by
    if limit > 0: conf.limit = limit

    msg = f"{conf.departure_country} -> {conf.arrival_country}, 日期: {conf.departure_date}, 时长小于 {conf.max_duration}小时, 总价小于 {conf.max_total_price}元, 排序规则: {conf.sort_by}, 返回: {conf.limit}条, 工作目录: {conf.work_dir}"
    write_to_file(f"{conf.work_dir}/TODO.md", f"## 查询两个国家之间的航班信息: \n", 'w')
    write_to_file(f"{conf.work_dir}/TODO.md", f"{msg}\n\n", 'a')
    write_to_file(f"{conf.work_dir}/TODO.md", f"1. 调用机票查询接口，查询航班数据并保存到 original_flight_data.md 文件中；\n", 'a')
    write_to_file(f"{conf.work_dir}/TODO.md", f"2. 根据 排序规则({conf.sort_by})、过滤规则(时长小于 {conf.max_duration}小时, 总价小于 {conf.max_total_price}元) 重新排序，排序后保存到 sorted_flight_data.md 文件中；\n", 'a')
    write_to_file(f"{conf.work_dir}/TODO.md", f"3. 挑出 {conf.limit}条 数据保存到 top_flight_data.md 文件中，并返回；\n", 'a')
    print(f"查询两个国家之间的航班信息: {msg}")

    # 查找两个国家之间的机票路线
    __find_international_flights_between_countries(conf)
    
    # 对航班数据进行排序
    title, content = __read_flights_data(f"{conf.work_dir}/original_flight_data.md")
    content, total_count, lose = __sort_flights_data(content, conf)

    # 写入排序后的航班数据到 sorted_flight_data.md 文件
    write_to_file(f"{conf.work_dir}/sorted_flight_data.md", f"{'\n'.join(title)}\n{'\n'.join(content)}\n", 'w')
    if lose > 0: write_to_file(f"{conf.work_dir}/sorted_flight_data.md", f"\n接口查询到的航班总数 {total_count} 条。其中 {lose} 条数据，不满足 飞行时长 < {conf.max_duration} 小时 或 总价 < {conf.max_total_price} 元 的限制。\n", 'a')
    
    # 写入前 20 条数据到 top_flight_data.md 文件
    write_to_file(f"{conf.work_dir}/top_flight_data.md", f"{'\n'.join(title)}\n{'\n'.join(content[:conf.limit])}\n", 'w')
    return f"{'\n'.join(title)}\n{'\n'.join(content[:conf.limit])}\n"


@mcp_tool()
def search_flights_between_cities(
    departure_list: list[str]=["北京"],
    arrival_list: list[str]=["上海"],
    departure_date: str="2026/02/10",
    max_duration: int=0,
    max_total_price: int=0,
    sort_by: list[str]=[],
    limit: int=0
):
    """
    可以查询两个城市之间的航班信息。
    
    你只要告诉我发城市列表和到达城市列表，我就能返回出发城市到到达城市之间的机票信息。

    Args:
        departure_list (list[str]): 出发城市列表，例如：["北京", "上海"]
        arrival_list (list[str]): 到达城市列表，例如：["广州", "纽约"]
        departure_date (str): 出发日期，格式为 YYYY-MM-DD
        max_duration (int): 最大飞行时长，单位：小时，默认值为 24小时
        max_total_price (int): 机票最大总价，单位：元，默认值为 5000元
        sort_by (list[str]): 排序规则字段，默认值为 ['total_price', 'duration'] 默认先按总价格排序，再按持续时间排序
        limit (int): 返回最优的记录条数，默认为20条
        
    Returns:
        str: 机票查询结果
    """
    conf = FlightConf() # 配置信息
    conf.work_dir = f"{os.getcwd()}/temp/{generate_timestamp_id()}" # 设置工作目录

    if len(departure_list) > 0: 
        conf.departure_city_list = departure_list
    else:
        return False, "出发城市列表 必填!"

    if len(arrival_list) > 0: 
        conf.arrival_city_list = arrival_list
    else:
        return False, "目的城市列表 必填!"

    if departure_date.strip(): 
        conf.departure_date = departure_date
    else:
        return False, "出发日期 必填!"

    if max_duration > 0: conf.max_duration = max_duration
    if max_total_price > 0: conf.max_total_price = max_total_price
    if len(sort_by) > 0: conf.sort_by = sort_by
    if limit > 0: conf.limit = limit

    msg = f"{conf.departure_city_list} -> {conf.arrival_city_list}, 日期: {conf.departure_date}, 时长小于 {conf.max_duration}小时, 总价小于 {conf.max_total_price}元, 排序规则: {conf.sort_by}, 返回: {conf.limit}条, 工作目录: {conf.work_dir}"
    write_to_file(f"{conf.work_dir}/TODO.md", f"## 查询两个城市之间的航班信息: \n", 'w')
    write_to_file(f"{conf.work_dir}/TODO.md", f"{msg}\n\n", 'a')
    write_to_file(f"{conf.work_dir}/TODO.md", f"1. 调用机票查询接口，查询航班数据并保存到 original_flight_data.md 文件中；\n", 'a')
    write_to_file(f"{conf.work_dir}/TODO.md", f"2. 根据 排序规则({conf.sort_by})、过滤规则(时长小于 {conf.max_duration}小时, 总价小于 {conf.max_total_price}元) 重新排序，排序后保存到 sorted_flight_data.md 文件中；\n", 'a')
    write_to_file(f"{conf.work_dir}/TODO.md", f"3. 挑出 {conf.limit}条 数据保存到 top_flight_data.md 文件中，并返回；\n", 'a')
    print(f"查询两个城市之间的航班信息: {msg}")

    # 查找出发城市列表和到达城市列表之间的机票路线
    __find_international_flights_between_city_lists(conf)
    
    # 对航班数据进行排序
    title, content = __read_flights_data(f"{conf.work_dir}/original_flight_data.md")
    content, total_count, lose = __sort_flights_data(content, conf)

    # 写入排序后的航班数据到 sorted_flight_data.md 文件
    write_to_file(f"{conf.work_dir}/sorted_flight_data.md", f"{'\n'.join(title)}\n{'\n'.join(content)}\n", 'w')
    if lose > 0: write_to_file(f"{conf.work_dir}/sorted_flight_data.md", f"\n接口查询到的航班总数 {total_count} 条。其中 {lose} 条数据，不满足 飞行时长 < {conf.max_duration} 小时 或 总价 < {conf.max_total_price} 元 的限制。\n", 'a')
    
    # 写入前 20 条数据到 top_flight_data.md 文件
    write_to_file(f"{conf.work_dir}/top_flight_data.md", f"{'\n'.join(title)}\n{'\n'.join(content[:conf.limit])}\n", 'w')
    return f"{'\n'.join(title)}\n{'\n'.join(content[:conf.limit])}\n"


def __search_international_one_way_by_fliggy(departure: str="北京", arrival: str="洛杉矶", departure_date: str="2026/02/10"):
    """国际单程机票价格查询"""

    print(f"国际单程机票价格查询: {departure} -> {arrival} -> {departure_date}")

    url = "https://sijipiao.fliggy.com/ie/flight_search_result_poller.do"

    headers = {
        "accept": "*/*",
        "accept-language": "zh-CN,zh;q=0.9",
        "cache-control": "no-cache",
        "cookie": "orderBy=totalPrice; lid=netlrs; wk_cookie2=1119937949fdb5bace00b975c928f73a; wk_unb=W8t8aeCZfog%3D; dnk=netlrs; tracknick=netlrs; _l_g_=Ug%3D%3D; unb=82880339; lgc=; cookie1=VqoJkBEC%2BLFTwOJFFWhqiypdK0PiRp9I8yzANtU5Kx4%3D; login=true; cookie17=W8t8aeCZfog%3D; cookie2=111cafb1b5748ae3bfe1620eafe62eeb; _nk_=netlrs; cancelledSubSites=empty; t=e630453d4ae01fdb1e39d85d28d82382; sg=s90; sn=; _tb_token_=7e3375a3b3573; XSRF-TOKEN=b7b3744a-164e-43d6-9dbf-0763a276ad1b; cna=UVg3H5RHfnMCAdz5NLpsvL3m; _m_h5_c=b184657ebc081915ef3fe19b00bd3ee8_1757576099899%3Beefe292eec4a34c91b95ac1198bdfefa; xlly_s=1; isg=BFZW-dTmIuhS5Bg1mn7SDNmQpwpY95ox9AXIP8C_FTn6g_YdKIR6Qb3wHx9vEpJJ; uc1=cookie14=UoYbw14evohq0g%3D%3D&existShop=false&cookie16=UIHiLt3xCS3yM2h4eKHS9lpEOw%3D%3D&cookie15=W5iHLLyFOGW7aA%3D%3D&cookie21=URm48syIYn73&pas=0; havana_lgc_exp=1789209616095; sgcookie=E100ciU3Rd3IYfqMWUaF%2Fv1aZ%2B2R%2BzkMQ0RHQlDWxO7PkI0Z%2FyeoCXUq7jnXtmnOyLYlhDp2vfcVcHeLbf829kUXREleLWk8jtC6u%2FTkt8r5K%2B0%3D; csg=c49ae700; mtop_partitioned_detect=1; _m_h5_tk=c2f0c0e8ecfa72db364e6f03c108e0d1_1758162185707; _m_h5_tk_enc=83bd852ba4858d83b586e115313acb8a; JSESSIONID=45F6956E7987A0206029C389B9B852FA; tfstk=gh9tSXN2UDE9YN4aXEcnnXjDd2onZXxaRF-7nZbgGeLpcFKDjtAGHEQprtxMIdDX9ETFsNOibeFv8FLMoZDNl-LDmdYmQs2Al3YQns9coKEwpUe0IVlwc-Bch40oEYxabtWjrsPnlIZNmi2_hrf_Ay6cDr6b7JKwbtrUxrGkOhrYCI9GftTfd6s52tsfcRiddijfhR1bCWaC8i6fh-1_dvsfq-a1lEtIvi7Clt16l9id0w_fhtG1ncQECgyv1_W6wlkQ3-2DpGCOk-bLHIvTE17WfwEmZpI9lat1J-wf7AKUcnTjk-XWaepANUDYLOxJO9s99mNC5_d6SM8xeR_6ph91Sdg3IN9wAQvDRmwfkBO5wQts-WIpOhvdtLg8H1RB7Q5vn2DPIQxk3_9sBz7haiLRCB3QBUszgLvR1gzlyiVsvDF4g1sFy1iNMc3sZkIdrc1_gS5JLMQoxZPbGrPcvamsnSPVMef..",
        "pragma": "no-cache",
        "referer": "https://sijipiao.fliggy.com/ie/flight_search_result.htm?searchBy=&b2g=0&formNo=-1&agentId=-1&needMemberPrice=true&searchJourney=%5B%7B%22depCityCode%22%3A%22BJS%22%2C%22arrCityCode%22%3A%22BKK%22%2C%22depCityName%22%3A%22%25E5%258C%2597%25E4%25BA%25AC%22%2C%22arrCityName%22%3A%22%25E6%259B%25BC%25E8%25B0%25B7%22%2C%22depDate%22%3A%222025-09-18%22%7D%5D&childPassengerNum=0&infantPassengerNum=0&tripType=0&cardId=",
        "sec-ch-ua": '"Not;A=Brand";v="99", "Google Chrome";v="139", "Chromium";v="139"',
        "sec-ch-ua-mobile": "?0",
        "sec-ch-ua-platform": '"macOS"',
        "sec-fetch-dest": "script",
        "sec-fetch-mode": "no-cors",
        "sec-fetch-site": "same-origin",
        "user-agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/139.0.0.0 Safari/537.36"
    }

    # 从原始URL中提取的参数
    params = {
        "_ksTS": f"{int(time.time() * 1000)}_{random.randint(100, 999)}",
        "callback": "jsonp728",
        "supportMultiTrip": "true",
        "searchBy": "",
        "childPassengerNum": "0",
        "infantPassengerNum": "0",
        "searchJourney": '[{"arrCityCode":"' + CityUtils.get_city_code(arrival) + '","arrCityName":"' + arrival + '","depCityCode":"' + CityUtils.get_city_code(departure) + '","depCityName":"' + departure + '","depDate":"' + departure_date + '","selectedFlights":[]}]',
        "tripType": "0",
        "searchCabinType": "0",
        "agentId": "-1",
        "controller": "1",
        "searchMode": "0",
        "b2g": "0",
        "formNo": "-1",
        "cardId": "",
        "needMemberPrice": "true"
    }

    try:
        while True:
            # 发送机票查询请求（需要请求两次，第一次价格较高）
            requests.get(url, headers=headers, params=params)
            time.sleep(random.randint(1, 3))
            response = requests.get(url, headers=headers, params=params)
            time.sleep(random.randint(8, 15))

            response.raise_for_status() # 检查响应状态码

            success, flight_info = FlightFliggyDataParse().parse(response.text) # 解析航班信息为 markdown 格式

            if success:
                break
            else:
                sleep_time = random.randint(180, 300)
                print(f"国际单程机票价格查询，接口返回失败，休眠 {sleep_time}S 后重试")
                time.sleep(sleep_time)
        
    except requests.exceptions.RequestException as e:
        return f"## 航班信息： 出发：{departure}， 到达：{arrival}， 出发日期：{departure_date}\nAn error occurred: {e}\n\n"

    print(f"国际单程机票价格查询，成功！")

    return f"## 航班信息： 出发：{departure}， 到达：{arrival}， 出发日期：{departure_date}\n{flight_info}\n\n"
    

def __find_international_flights_between_countries(conf: FlightConf=None):
    """
    查找两个国家之间的机票路线，查找到的数据保存到 original_flight_data.md 文件中
    
    Args:
        conf (FlightConf): 配置对象
    
    Returns:
        bool: 查找成功返回 True，否则返回 False
    """
    conf.departure_city_list = CityUtils.get_country_city_list(conf.departure_country) # 出发城市列表
    conf.arrival_city_list = CityUtils.get_country_city_list(conf.arrival_country) # 到达城市列表

    return __find_international_flights_between_city_lists(conf)


def __find_international_flights_between_city_lists(conf: FlightConf=None):
    """
    查找出发城市列表和到达城市列表之间的机票路线，查找到的数据保存到 original_flight_data.md 文件中
    
    Args:
        conf (FlightConf): 配置对象
    
    Returns:
        bool: 查找成功返回 True，否则返回 False
    """

    i = 0
    for departure in conf.departure_city_list:
        for arrival in conf.arrival_city_list:
            i = i + 1
            print(f"循环第 {i} 次， {departure} -> {arrival}")
            write_to_file(f"{conf.work_dir}/original_flight_data.md", __search_international_one_way_by_fliggy(departure, arrival, conf.departure_date), 'a')
    
    return True

def __sort_flights_data(flights_data, conf: FlightConf=None):
    """
    对航班数据进行排序，根据总价格和持续时间
    
    Args:
        flights_data (list[str]): 航班数据列表，每个元素是一个航班的 markdown 格式字符串
        conf (FlightConf): 配置对象
    
    Returns:
        list[str]: 排序后的航班数据列表
        len(df_sorted): 待排序的航班数据条数
        lose(int): 过滤掉的航班数量
    """

    # 解析成为表格数据
    table_data = []
    for line in flights_data:
        if line.strip() and '|' in line:
            
            # 分割并清理所有字段
            parts = line.split('|') # 先按 | 分割行数据
            if line.startswith('|'): # 移除可能的首空元素（由于 | 在行首）
                parts = parts[1:]
            if line.endswith('|'): # 移除可能的尾空元素（由于 | 在行尾）
                parts = parts[:-1]
            cells = [part.strip() for part in parts] # 将分隔后的数据转为数组
            if len(cells) >= 14:  # 确保有14列数据
                table_data.append([cells[0], cells[3], cells[4]])
    
    # 转换为 DataFrame
    df = pd.DataFrame(table_data, columns=['id', 'duration', 'total_price'])

    # 清理和转换数据
    df['total_price'] = df['total_price'].str.replace('¥', '').astype(float)
    df['duration'] = df['duration'].str.replace('小时', '').str.strip().astype(float)
    
    # 排序
    df_sorted = df.sort_values(by=conf.sort_by)

    # 根据排序后的结果生成最终结果数据
    result = []
    lose = 0 # 过滤掉的航班数量
    for index, row in df_sorted.iterrows():
        if row['duration'] > conf.max_duration or row['total_price'] > conf.max_total_price: # 过滤掉时间大于 max_duration 和价格大于 max_total_price 的数据
            lose = lose + 1
            continue
        result.append(flights_data[index])

    return result, len(df_sorted), lose


def __read_flights_data(filename: str):
    """
    读取文件中的航班数据
    
    Args:
        filename (str): 文件名，包含文件路径
    
    Returns:
        tuple: 包含标题行和内容数据的元组 (title, content)
    """
    is_title = False # 是否读到标题
    title = [] # 标题行
    content = [] # 内容数据
    with open(filename, 'r', encoding='utf-8') as file:
        for line in file:
            line_str = line.strip() # 读入一行数据

            if line_str.startswith('| ID |'): # 头
                if not is_title:
                    title.append(line_str) 
            elif line_str.startswith('|---|'): # 头
                if not is_title:
                    title.append(line_str)
                    is_title = True
            elif line_str.startswith('|'): # 航班数据
                content.append(line_str)

    return title, content

# 调用 AI 处理航班数据，暂时没用注释
# def read():
#     ai_utils = AIUtils()
#     with open('example2.txt', 'r', encoding='utf-8') as file:
#         content = []
#         for line in file:
#             line_str = line.strip() # 读入一行数据

#             if line_str.startswith('## 航班信息：'): # 头
#                 title = line_str 
#             elif not line_str: # 尾
#                 if len(content) > 0:
#                     result = ai_utils.process_query(flight_optimizer_prompt.replace('{data}', '\n'.join(content))) # 调用 AI 返回最优航班 ID

#                     # 根据 AI 返回的航班 ID，从航班数据中查出
#                     best_data = []
#                     for id_str in result.split(','):
#                         if id_str.strip():
#                             for ss in content:
#                                 if id_str in ss:
#                                     best_data.append(ss)
                    
#                     if len(best_data) > 0:               
#                         write_to_file('zlh.txt', f"{title}\n{'\n'.join(best_data)}\n\n", 'a')

#                     content = [] # 清空列表，准备下一个标题下的内容

#                     # return
#             else: # 航班数据
#                 content.append(line_str)
#     return
